#include <stdlib.h>
#include <iostream>
#include <conio.h>
#include <string>
#include "OpenNI.h"
#include "opencv2/core/core.hpp"
#include "opencv2/highgui/highgui.hpp"
#include "opencv2/imgproc/imgproc.hpp"

#define RESOLUTION 640,480
#define RECORDRESOLUTION 590,440
#define ADRESOLUTION 45,40,590,440
#define FPS 20
#define GRAYTH 10
#define REPAIRERANGE 5
#define COLORTH 10

using namespace std;
using namespace cv;
using namespace openni;

//Openni status
Status result = STATUS_OK;
//open device
Device device;
//OpenNI2 image
VideoFrameRef oniDepthImg;
VideoFrameRef oniColorImg;
//create stream
VideoStream oniDepthStream;
VideoStream oniColorStream;
//set video mode
VideoMode modeDepth;
VideoMode modeColor;
//OpenCV image
cv::Mat cvDepthImg;
cv::Mat cvDepthImg2;
cv::Mat cvColorImg;
cv::Mat cvColorImg2;
//OpenCV adjusted image
cv::Mat cvAdjustDepthImg;
cv::Mat cvAdjustColorImg;
//Resolution
Size se=Size(RESOLUTION);
Size recordse=Size(RECORDRESOLUTION);

void CheckOpenNIError( Status result, string status )
{
    if( result != STATUS_OK )
        cerr << status << " Error: " << OpenNI::getExtendedError() << endl;
}
void iniOpenNI(void );
void releaseResource(void );
void RemoveNoiseRecord(void );
void RemoveNoise(void );
bool pixelRepaire(int ,int ,int );
bool rangeRepaire(int ,int ,int );


int main( int argc, char** argv )
{
    iniOpenNI();
    RemoveNoiseRecord();
    releaseResource();
    return 0;
}

void iniOpenNI()
{
    // initialize OpenNI2
    result = OpenNI::initialize();
    CheckOpenNIError( result, "initialize context" );

    //open video device
    result = device.open( openni::ANY_DEVICE );
    CheckOpenNIError( result, "initialize context" );

    //creat depth stream
    result = oniDepthStream.create( device, openni::SENSOR_DEPTH );
    CheckOpenNIError( result, "initialize context" );
    //set depth mode
    modeDepth.setResolution( RESOLUTION );
    modeDepth.setFps( FPS );
    modeDepth.setPixelFormat( PIXEL_FORMAT_DEPTH_1_MM );
    oniDepthStream.setVideoMode(modeDepth);
    // start depth stream
    result = oniDepthStream.start();
    CheckOpenNIError( result, "initialize context" );

    // create color stream
    result = oniColorStream.create( device, openni::SENSOR_COLOR );
    CheckOpenNIError( result, "initialize context" );
    // set color video mode
    modeColor.setResolution( RESOLUTION );
    modeColor.setFps( FPS );
    modeColor.setPixelFormat( PIXEL_FORMAT_RGB888 );
    oniColorStream.setVideoMode( modeColor);
    // start color stream
    result = oniColorStream.start();
    CheckOpenNIError( result, "initialize context" );

    // set depth and color imge registration mode
    if( device.isImageRegistrationModeSupported(IMAGE_REGISTRATION_DEPTH_TO_COLOR ) )
    {
        cout << "support" << endl;
        device.setImageRegistrationMode( IMAGE_REGISTRATION_DEPTH_TO_COLOR );
    }

}

void releaseResource()
{
    //OpenNI2 destroy
    oniDepthStream.destroy();
    oniColorStream.destroy();
    device.close();
    OpenNI::shutdown();
}



void RemoveNoiseRecord()
{
    std::cout<<"\t\t\t------RemoveNoiseRecord-------"<<std::endl;

    char *DepthFilename = "oringindepthvideo.avi";
    char *ColorFilename = "oringincolorvideo.avi";
    char *removeDepthFilename = "removedepthvideo.avi";
    char *removeColorFilename = "removecolorvideo.avi";


    VideoWriter removecolorVideoWriter=VideoWriter(removeColorFilename,CV_FOURCC('X','V','I','D'),FPS,recordse);
    VideoWriter removedepthVideoWriter=VideoWriter(removeDepthFilename,CV_FOURCC('X','V','I','D'),FPS,recordse);
    VideoWriter colorVideoWriter=VideoWriter(ColorFilename,CV_FOURCC('X','V','I','D'),FPS,recordse);
    VideoWriter depthVideoWriter=VideoWriter(DepthFilename,CV_FOURCC('X','V','I','D'),FPS,recordse);
    //    VideoWriter colorVideoWriter=VideoWriter(ColorFilename,CV_FOURCC('X','V','I','D'),FPS,se);
    //    VideoWriter depthVideoWriter=VideoWriter(DepthFilename,CV_FOURCC('X','V','I','D'),FPS,se);
    namedWindow("after-gray",1);
    namedWindow("after-depth",1);
    namedWindow("befor-color",1);
    namedWindow("befor-depth",1);
    while(true)
    {
        if( oniColorStream.readFrame( &oniColorImg ) == STATUS_OK )
        {
            // convert data into OpenCV type
            cv::Mat cvRGBImg( oniColorImg.getHeight(), oniColorImg.getWidth(), CV_8UC3, (void*)oniColorImg.getData() );
            cv::cvtColor( cvRGBImg, cvColorImg2, CV_RGB2BGR );
            cvColorImg2=Mat(cvColorImg2,Rect(ADRESOLUTION));
            colorVideoWriter.write(cvColorImg2);
            cv::imshow("befor-color", cvColorImg2 );
            cvtColor(cvRGBImg,cvColorImg,CV_RGB2GRAY);
            //colorVideoWriter.write(cvColorImg);
        }
        if( oniDepthStream.readFrame( &oniDepthImg ) == STATUS_OK )
        {
            cv::Mat cvRawImg16U( oniDepthImg.getHeight(), oniDepthImg.getWidth(), CV_16UC1, (void*)oniDepthImg.getData() );
            cvRawImg16U.convertTo( cvDepthImg, CV_8UC1, 255.0/(oniDepthStream.getMaxPixelValue()));
            cv::cvtColor(cvDepthImg,cvDepthImg2,CV_GRAY2BGR);
            cvDepthImg2=Mat(cvDepthImg2,Rect(ADRESOLUTION));
            depthVideoWriter.write(cvDepthImg2);
            cv::imshow( "befor-depth", cvDepthImg2 );
        }

        cvAdjustDepthImg=Mat(cvDepthImg,Rect(ADRESOLUTION));
        cvAdjustColorImg=Mat(cvColorImg,Rect(ADRESOLUTION));

        RemoveNoise();

        cvtColor(cvAdjustColorImg,cvAdjustColorImg,CV_GRAY2BGR);
        cvtColor(cvAdjustDepthImg,cvAdjustDepthImg,CV_GRAY2BGR);

        removecolorVideoWriter.write(cvAdjustColorImg);
        removedepthVideoWriter.write(cvAdjustDepthImg);
        std::cout<<"removecolorVideoWriter"<<std::endl;
        std::cout<<"\tremovedepthVideoWriter"<<std::endl;

        imshow("after-gray",cvAdjustColorImg);
        imshow("after-depth",cvAdjustDepthImg);

        int key;
        key=waitKey(10);
        if(key==27)
        {
            break;
        }
    }
    removecolorVideoWriter.release ();
    removedepthVideoWriter.release ();
    colorVideoWriter.release ();
    depthVideoWriter.release ();

    destroyWindow("after-gray");
    destroyWindow("after-depth");
    destroyWindow("befor-color");
    destroyWindow("befor-depth");
}

void RemoveNoise()
{
    clock_t start,finish;
    double totaltime=0.0;
    start=clock();
    for(int j=(cvAdjustDepthImg.rows-1);j>=0;j--)//depthImage.rows
    {
        const uchar* mj=cvAdjustDepthImg.ptr<uchar>(j);
        for(int i=(cvAdjustDepthImg.cols-1);i>=0;i--)//depthImage.cols
        {
            if(mj[i]<=GRAYTH)
            {
                uchar colorpixel=cvAdjustColorImg.at<uchar>(j,i);
                bool reResult=false;
                if(colorpixel<GRAYTH*5)
                {
                    for(int k=1;k<REPAIRERANGE*3;k++)
                    {
                        reResult=pixelRepaire(i,j,k);
                        if(reResult)
                            break;
                    }
                    //go down
                    if(!reResult)
                    {
                        for(int k=1;k<=30;k++)
                        {
                            if((j+k)<440)
                            {
                                if(cvAdjustDepthImg.at<uchar>(j+k,i)>GRAYTH)
                                {
                                    cvAdjustDepthImg.at<uchar>(j,i)=cvAdjustDepthImg.at<uchar>(j+k,i);
                                    reResult=true;
                                    break;
                                }
                            }
                            else
                            {
                                break;
                            }
                        }
                    }
                    //go up
                    if(!reResult)
                    {
                        for(int k=1;k<=30;k++)
                        {
                            if((j-k)>=0)
                            {
                                if(cvAdjustDepthImg.at<uchar>(j-k,i)>GRAYTH)
                                {
                                    cvAdjustDepthImg.at<uchar>(j,i)=cvAdjustDepthImg.at<uchar>(j-k,i);
                                    reResult=true;
                                    break;
                                }
                            }
                            else
                            {
                                break;
                            }
                        }
                    }
                }
                else
                {
                    for(int k=1;k<30;k++)
                    {
                        if((i+k)<590 && !reResult)
                        {
                            if(abs(cvAdjustColorImg.at<uchar>(j,i+k)-colorpixel)<=COLORTH && cvAdjustDepthImg.at<uchar>(j,i+k)>GRAYTH)
                            {
                                cvAdjustDepthImg.at<uchar>(j,i)=cvAdjustDepthImg.at<uchar>(j,i+k);
                                reResult=true;
                            }
                        }
                        else
                        {
                            break;
                        }
                    }
                }
                if(!reResult)
                {
                    for(int k=1;k<REPAIRERANGE;k++)
                    {
                        reResult=pixelRepaire(i,j,k);
                        if(reResult)
                            break;
                    }
                }
                if(!reResult)
                {
                    for(int k=0;k<REPAIRERANGE*3;k++)
                    {
                        reResult=rangeRepaire(i,j,k);
                        if(reResult)
                            break;
                    }
                }
            }
        }
    }
    finish=clock();
    totaltime=(double)(finish-start)/CLOCKS_PER_SEC;
}

bool pixelRepaire(int i,int j,int repaireRange)
{
    uchar colorpixel=cvAdjustColorImg.at<uchar>(j,i);
    int x=0;
    int y=0;
    int n=0;
    int sum=0;
    for(y=j-repaireRange;y<=j+repaireRange;y++)
    {
        if(y>=0 && y<440)
        {

            if(y==(j-repaireRange) ||  y==(j+repaireRange))
            {
                for(x=i-repaireRange;x<=i+repaireRange;x++)
                {
                    if(x>=0 && x<590)
                    {
                        if(abs(cvAdjustColorImg.at<uchar>(y,x)-colorpixel)<=COLORTH && cvAdjustDepthImg.at<uchar>(y,x)>GRAYTH)
                        {
                            n++;
                            sum=sum+cvAdjustDepthImg.at<uchar>(y,x);
                        }
                    }

                }
            }

            else
            {
                x=i-repaireRange;
                if(x>=0 && x<590)
                {
                    if(abs(cvAdjustColorImg.at<uchar>(y,x)-colorpixel)<=COLORTH && cvAdjustDepthImg.at<uchar>(y,x)>GRAYTH)
                    {
                        n++;
                        sum=sum+cvAdjustDepthImg.at<uchar>(y,x);
                    }
                }
                x=i+repaireRange;
                if(x>=0 && x<590)
                {
                    if(abs(cvAdjustColorImg.at<uchar>(y,x)-colorpixel)<=COLORTH && cvAdjustDepthImg.at<uchar>(y,x)>GRAYTH)
                    {
                        n++;
                        sum=sum+cvAdjustDepthImg.at<uchar>(y,x);
                    }
                }
            }
        }
    }
    if(n<repaireRange*2)
    {return false;}
    else
    {
        cvAdjustDepthImg.at<uchar>(j,i)=(uchar)(sum/n);
        return true;
    }

}

bool rangeRepaire(int i,int j,int repaireRange)
{
    uchar colorpixel=cvAdjustColorImg.at<uchar>(j,i);
    int x=0;
    int y=0;
    int n=0;
    int sum=0;
    for(y=j-repaireRange;y<=j+repaireRange;y++)
    {
        if(y>=0 && y<440)
        {
            for(x=i-repaireRange;x<=i+repaireRange;x++)
            {
                if(x>=0 && x<590)
                {
                    if(cvAdjustDepthImg.at<uchar>(y,x)>GRAYTH)
                    {
                        n++;
                        sum=sum+cvAdjustDepthImg.at<uchar>(y,x);
                    }
                }
            }
        }
    }
    if(n<=repaireRange*2)
    {
        return false;
    }
    else
    {
        cvAdjustDepthImg.at<uchar>(j,i)=(uchar)(sum/n);
        return true;
    }
}


